"""
将 MIDOG 2021 中  tif 转换成 yolo 系列需要的数据集

https://midog.deepmicroscopy.org/download-dataset/

m21
[1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149]
[2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]
[3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

f1:
    train:
        [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

    test:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149]

f2:
    train:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

    test:
        [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]

f3:
    train:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149
        ,2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]

    test:
        [3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]


"""
import json
import os
import numpy as np
from PIL import Image
import cv2
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor

wsi_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/tiff_sn/"
midog_json = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/MIDOG.json"
data = json.load(open(midog_json))
png_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/640_pn_/images/"
label_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/640_pn_/labels/"
os.makedirs(png_path, exist_ok=True)
os.makedirs(label_path, exist_ok=True)

patch_size = 640  # 训练 image 大小
react_size = 50  # 框的大小 40X 下
react_size_percent = react_size / patch_size

filter_arr = [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

save_gt = False
if save_gt:
    gt_path = f"/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/640_pn/gt/"
    os.makedirs(gt_path, exist_ok=True)

def do_convert(wsi_file):
    image_id = int(wsi_file[:3])
    if image_id not in filter_arr:
        return
    file_name = wsi_file.split('.')[0]
    file_path = os.path.join(wsi_path, wsi_file)
    img = cv2.imread(file_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    height, width, _ = img.shape
    mask = np.zeros((height, width)).astype(np.uint8)

    annotations = [anno for anno in data["annotations"] if
                   anno["image_id"] == image_id and anno["category_id"] == 1]

    if len(annotations) > 0:
        for anno in annotations:
            bbox_arr = anno["bbox"]
            x = bbox_arr[1]
            y = bbox_arr[0]
            x = int(x + 25)
            y = int(y + 25)
            mask[x, y] = 1

    overlap = 0
    for h in range(0, height, patch_size - overlap):
        h_start = h
        h_end = h + patch_size
        if h_end > height:
            h_start = height - patch_size
            h_end = height
        for w in range(0, width, patch_size - overlap):
            w_start = w
            w_end = w + patch_size
            if w_end > width:
                w_start = width - patch_size
                w_end = width

            # 切割块
            png_patch = img[h_start:h_end, w_start:w_end]
            mask_patch = mask[h_start:h_end, w_start:w_end]
            gt_patch = None

            png_name = f"{file_name}_{w_start}_{h_start}.png"
            txt_name = f"{file_name}_{w_start}_{h_start}.txt"

            if mask_patch.max() == 0:
                Image.fromarray(png_patch).save(png_path + png_name)
            else:
                if save_gt:
                    gt_patch = np.asarray(png_patch)
                points = np.argwhere(mask_patch == 1)
                for point in points:
                    label = (
                        0,
                        point[1] / patch_size,
                        point[0] / patch_size,
                        react_size_percent,
                        react_size_percent
                    )
                    with open(label_path + txt_name, 'a') as f:
                        f.write(('%g ' * len(label)).rstrip() % label + '\n')
                Image.fromarray(png_patch).save(png_path + png_name)
                if save_gt:
                    Image.fromarray(gt_patch).save(gt_path + png_name)
    print(f"{wsi_file} done")


# with ThreadPoolExecutor(max_workers=20) as executor:
#     for wsi_file in os.listdir(wsi_path):
#         executor.submit(do_convert, wsi_file)


for wsi_file in sorted(os.listdir(wsi_path)):
    do_convert(wsi_file)